package middleware

type Location = byte

const (
	HeaderLocation Location = iota
	CookieLocation
)

func TrackWrap(key string, location Location, locationKey string) Handler {
	var handler func(key string, context Context) string
	switch location {
	case HeaderLocation:
		handler = headerLocation
	case CookieLocation:
		handler = cookieLocation
	}
	return func(context Context, next Next) error {
		v := handler(locationKey, context)
		if v != "" {
			// 设定log标识
			context.SetLogger(context.LoggerWrap(key, v))
			context.SetID(v)
		}
		return next(nil)
	}
}

func headerLocation(key string, context Context) string {
	return context.Header(key)
}

func cookieLocation(key string, context Context) string {
	v, err := context.Cookie(key)
	if err != nil {
		return ""
	} else {
		return v.Value
	}
}
